import os
import warnings
import multiprocessing as mp

import tqdm
import librosa

import numpy as np
import pandas as pd

import torch.utils.data as td

from typing import Any
from typing import Dict
from typing import List
from typing import Tuple
from typing import Union
from typing import Optional

from model.mvclip import MusicVideoCLIP

#######新增的#######
import simplejpeg

##############################################
#因为模型要提出音频需要调用class mvclip
#可考虑将其函数分离出class
##############################################
MODEL_FILENAME = 'AudioCLIP-Full-Training.pt'
mvclip = MusicVideoCLIP(pretrained=f'../assets/{MODEL_FILENAME}')


####################################################
#EmoMV数据集需要修改的地方，提取音频，光流特征，文本特征
#说明 此处新增加参数：transform_video
####################################################

class EmoMV(td.Dataset):
    def __init__(self,
                 root: str,
                 sample_rate: int = 22050,
                 train: bool = True,
                 fold: Optional[int] = None,
                 transform_music=None,
                 transform_video=None,
                 target_transform=None,
                 **_):

        super(EmoMV, self).__init__()

        self.sample_rate = sample_rate
        
        #####################esc50.csv####################
        #这个文件里面存在text信息
        #
        ##################################################
        meta = self.load_meta(os.path.join(root, 'meta', 'emomv.csv'))
        
        
        #####################################################
        #fold 是指交叉验证折数 默认5折 也对应其meta文件的filename
        #
        #####################################################
        if fold is None:
            fold = 5

        self.folds_to_load = set(meta['fold'])

        if fold not in self.folds_to_load:
            raise ValueError(f'fold {fold} does not exist')

        self.train = train
        
        ##############数据增强处理#############
        self.transform_music = transform_music
        self.transform_video = transform_video
        
        
        

        if self.train:
            self.folds_to_load -= {fold}
        else:
            self.folds_to_load -= self.folds_to_load - {fold}
        
        ############################################
        #下列 self.data为注解
        #
        ############################################
        self.data: Dict[Union[str, int], Dict[str, Any]] = dict()
#         self.load_data(meta, os.path.join(root, 'audio')) # 具体指到audio的文件夹
        self.load_data(meta,os.path.join(root, 'video'))
        
        ############################################
        # self.indices:索引数据
        #
        ############################################
        self.indices = list(self.data.keys())  # 用于确认数据集长度的，以便后续的索引
        
        
        
        ############################################
        # 将类进行编码：具体就是目标和类之间的编码
        #
        ############################################
        self.class_idx_to_label = dict() # 用于编码text
        for row in self.data.values():
            idx = row['target']
            label = row['category']
            self.class_idx_to_label[idx] = label
        self.label_to_class_idx = {lb: idx for idx, lb in self.class_idx_to_label.items()}

        self.target_transform = target_transform

    ###############################################
    #@staticmethod装饰器可以将一个方法转换为静态方法
    #可以直接调用，而无需创建类实体
    #
    ###############################################
    
    @staticmethod
    def load_meta(path_to_csv: str) -> pd.DataFrame:
        
        meta = pd.read_csv(path_to_csv)

        return meta

    
    ##############################重点修改处，此处是对音频的提取##################################
    #具体就是对一个wav文件进行修改
    #需要返回wav梅尔滤波图
    #注意此处就需要引入mvclip，可以考虑将函数脱离出来 
    #同时注意数据集成了batch
    #filename: e.g. ESC-50/audio/3-144259-A-29.wav   
    ###########################################################################################
    @staticmethod
    def _load_worker(idx: int, filename: str, sample_rate: Optional[int] = None) -> Tuple[int, int, np.ndarray]:
        """
        注意：此加载的是一个wav文件，如此，若已经分割好了，需要对meta的csv文件进行修改，单位是segment；
        若没有分割好，就需要在此处完成分割任务，并进行加载，注意此时返回的是一个wav的数据，单位是wav；
        """
        
        ###########原始的############
#         wav, sample_rate = librosa.load(filename, sr=sample_rate, mono=True)
        ############################
        
        
        ####################新增加的，利用原始训练过的audio-head(ESResNeXt)###################
        # 可参考的链接 https://blog.csdn.net/chumingqian/article/details/123404790
        # filename: e.g. ESC-50/audio/3-144259-A-29.wav
        ###################################################################################
        
#         wav,sample_rate = librosa.load(filename, sr=sample_rate, mono=True)
#         spec = mvclip.audio.spectrogram(torch.from_numpy(wav.reshape(1, 1, -1)))
#         spec = np.ascontiguousarray(spec.numpy()).view(np.complex64)
#         pow_spec = 10 * np.log10(np.abs(spec) ** 2 + 1e-18).squeeze()
#         track_cut = np.zeros((44100,))
#         length = min(len(track), 44100)
#         track_cut[:length] = track[:length]

#         ############原始的###########
# #         if wav.ndim == 1:
# #             wav = wav[:, np.newaxis]

# #         wav = wav.T * 32768.0
#         ############################
        
        
#         ########一wav中的一片段#########
#         #需要将seg_music转为np.float类型
#         #
#         ###############################
        
        # seg_music=(track_cut,pow_spec)
        
        
        track, b_ = librosa.load(filename, sr=sample_rate, dtype=np.float32)
        track6=track.copy()
        
        ##########音频分成6段###########
        tracklen=list(track6.shape)[0]//6*6
        track6=track6[:tracklen].reshape(6,-1)

        # compute spectrograms using trained audio-head (fbsp-layer of ESResNeXt)
        # thus, the actual time-frequency representation will be visualized
        spec = mvclip.audio.spectrogram(torch.from_numpy(track.reshape(1, 1, -1)))
        spec = np.ascontiguousarray(spec.numpy()).view(np.complex64)
        pow_spec = 10 * np.log10(np.abs(spec) ** 2 + 1e-18).squeeze()
        
        #########音频信息对齐##########
        powlenrow=list(pow_spec.shape)[0]
        powlendowm = list(pow_spec.shape)[1]
        pow_spec_cut = powlenrow//6*6
        pow_spec6 = pow_spec[:pow_spec_cut] 
        
        ##########分成6段#########
        pow_spec6 = pow_spec6.reshape(6,-1,powlendowm)
        len6=list(track6.shape)[0] # 其实就是6
        music_single6=[(track6[i],pow_spec6[i])for i in range(len6)]


        ######################jpg filename##########################
        video_filename=os.path.basename(filename).split('.')[0] # 3-144259-A-29
        prefix=filename.split('.')[0] # ESC-50
        path2video='/'.join([prefix,'video',video_filename]) #'ESC-50/video/3-144259-A-29'
        
        
        ############################新增加的，对video特征的提取###############################
        # 具体来说就是光流图和rgb图
        # 此处需要定义好文件名和对应的 光流图的文件夹名，光流图的提取需要提前的环境直接提取，
        # 有部分论文即采用这种方式
        # 拟定采用以下方式：
        # 一个wav有6段wav，
        # 结构采用一个video：[rgb,flowx,flowy],rgb :[[3,224,224]...]
        # 一个batch为[video,video,...] 分段在前向完成
        ####################################################################################
        
        
        ################rgb##################
        video_rgb_single = list()
        for idx, f in enumerate(sorted(glob.glob(path2video+'/img*.jpg'))):
            with open(f, 'rb') as jpg:
                image = simplejpeg.decode_jpeg(jpg.read())
                video_rgb_single.append(image)
        
        ################flow_x##################
        video_flow_x_single = list()
        for idx, f in enumerate(sorted(glob.glob(path2video+'/flow_x*.jpg'))):
            with open(f, 'rb') as jpg:
                image = simplejpeg.decode_jpeg(jpg.read())
                video_flow_x_single.append(image)
        
        ################flow_y##################
        video_flow_y_single = list()
        for idx, f in enumerate(sorted(glob.glob(path2video+'/flow_y*.jpg'))):
            with open(f, 'rb') as jpg:
                image = simplejpeg.decode_jpeg(jpg.read())
                video_flow_y_single.append(image)

        video6 = [video_rgb_single, video_flow_x_single, video_flow_y_single] #5维度，其中如video_rgb 4维 [[3,224,224],]
        
        
#         return idx, sample_rate, wav.astype(np.float32)
        return idx, sample_rate, music_single6,video6

    
    
    
    def load_data(self, meta: pd.DataFrame, base_path: str):
        items_to_load = dict()
        
        ########################文件的读取###########################
        #首先加载到items_to_load，加载的是wav文件的读取
        #
        ############################################################
        
        ##############csv文件的读取##############
        for idx, row in meta.iterrows():
            if row['fold'] in self.folds_to_load:
                #命名这里，csv里是带.wav的
                items_to_load[idx] = os.path.join(base_path, row['Filename']), self.sample_rate
        
        
        items_to_load = [(idx, path, sample_rate) for idx, (path, sample_rate) in items_to_load.items()]
        
        num_processes = os.cpu_count()
        warnings.filterwarnings('ignore')
        with mp.Pool(processes=num_processes) as pool:
            tqdm.tqdm.write(f'Loading {self.__class__.__name__} (train={self.train})')
            
            # 需要对应1段music(含6段) 代码实现切分
            ####################此处进行video特征的增加#######################
            for idx, sample_rate, music ,video in pool.starmap(
                    func=self._load_worker,  # 具体就是对一个wav的提取
                    iterable=items_to_load,
                    chunksize=int(np.ceil(len(items_to_load) / num_processes)) or 1
            ):
                row = meta.loc[idx]
            
                self.data[idx] = {
                    'music': music,       # 注意此处的数据为tuple类型，非np.float32格式，同时注意：命名audio->music
                    'video': video,       # 新增的video特征
                    'sample_rate': sample_rate,
                    'target': row['target'],
                    'category': row['category'].replace('_', ' '),
                    'fold': row['fold'],
                    # 'esc10': row['esc10'] # The esc10 column indicates if a given file belongs to the ESC-10 subset
                }

    
    ##############################################################################################
    #提取batch的核心修改处，需要在此处完成对数据集加载，之后torchvision直接shuffle即可
    #注意此处是对audio的处理，返回audio以及 target，对应关系为一个文件，因此此处通过stack的形式进行修改
    #return : music,video,[target] #target 为情绪值,需要注意的是情绪值在数据集meta中的csv文件中进行定义
    #另外需要注意的是：category对应情绪值，对于图像对应情绪可能会存在问题,即为效果不好
    ##############################################################################################
    
    def __getitem__(self, index: int) -> Tuple[np.ndarray, Optional[np.ndarray], List[str]]:
        if not (0 <= index < len(self)):
            raise IndexError

        music: np.ndarray = self.data[self.indices[index]]['music']
        video: list = self.data[self.indices[index]]['video']
        target: str = self.data[self.indices[index]]['category']
        
        if self.transform_music is not None:
            music = self.transform_music(music)
        
        ######################添加video图像的transform方法#################
        if self.video_transform is not None:
            for j in range(len(video)): # 3
                for i in range(len(video[j])): 
                    video_single = torch.stack([image_transforms(image) for image in video[j][i]])
                    video[j][i] = video_single

        
        ##########文本的transform##########
        if self.target_transform is not None:
            target = self.target_transform(target)


        return music, video, [target]

    def __len__(self) -> int:
        return len(self.indices)
    
    
    
